"""
NMS简化代码


"""

import matplotlib.pyplot as plt
import numpy as np


def nms(boxes, thresh=0.7, is_show_demo=False):
    """
    非极大值抑制

    提示：需要基础：Numpy广播和切片、使用了一点点python闭包
    提示：演示时：保留：红色、抑制：蓝色、未定：黑色。
    提示：这里有一些输出，是为了让调用者了解原理而设置，正式使用请注释掉。

    :param boxes: 输入矩阵，每行5个数据，依次是：x1, y1, x2, y2, 置信度
    :param thresh: IOU小于thresh的不抑制
    :param is_show_demo: 是否用画图库进行演示
    :return: 抑制后的矩阵，格式同boxes
    """

    def sep(label='', cnt=32):
        print('-' * cnt, label, '-' * cnt, sep='')

    def draw_boxes(boxes, idx_arr, color):
        x1 = boxes[idx_arr, 0]
        y1 = boxes[idx_arr, 1]
        x2 = boxes[idx_arr, 2]
        y2 = boxes[idx_arr, 3]
        plt.plot([x1, x2], [y1, y1], color=color)
        plt.plot([x1, x2], [y2, y2], color=color)
        plt.plot([x1, x1], [y1, y2], color=color)
        plt.plot([x2, x2], [y1, y2], color=color)
        for idx in idx_arr:
            plt.annotate(f'#{idx}({boxes[idx, 4]})', xy=[boxes[idx, 0], boxes[idx, 1]])

    def draw_boxes_groups(boxes, idx_arr_to_keep, idx_arr_to_remove, idx_arr_uncertain, title):
        nonlocal spn
        spn += 1
        plt.subplot(spr, spc, spn)
        plt.title(title)
        draw_boxes(boxes, idx_arr_uncertain, 'k')
        draw_boxes(boxes, idx_arr_to_remove, 'b')
        draw_boxes(boxes, idx_arr_to_keep, 'r')
        plt.ylim(plt.ylim()[::-1])

    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    idx_ordered_by_c = scores.argsort()[::-1]
    print(idx_ordered_by_c)

    idx_arr_to_keep = np.array([], dtype=np.int32)
    idx_arr_to_remove = np.array([], dtype=np.int32)
    idx_arr_uncertain = idx_ordered_by_c.copy()
    EPS = 1e-10

    no_of_action = 0
    while len(idx_arr_uncertain) > 0:
        no_of_action += 1
        sep(no_of_action)

        #########################################################
        # select top C
        idx_to_keep = idx_arr_uncertain[0]

        #########################################################
        # kepp it and it is never uncertain from now on
        idx_arr_to_keep = np.append(idx_arr_to_keep, [idx_to_keep])
        idx_arr_to_remove = np.delete(idx_arr_to_remove, range(len(idx_arr_to_remove)))
        idx_arr_uncertain = idx_arr_uncertain[1:]
        print('before')
        print('idx_to_keep', idx_arr_to_keep)
        print('idx_arr_to_remove', idx_arr_to_remove)
        print('idx_arr_uncertain', idx_arr_uncertain)

        #########################################################
        # plot condition of before
        if is_show_demo:
            plt.figure(figsize=[12, 6])
            spr = 1
            spc = 2
            spn = 0
            draw_boxes_groups(boxes, idx_arr_to_keep, idx_arr_to_remove, idx_arr_uncertain, f'#{no_of_action} before')

        #########################################################
        # calculate IOU
        if len(idx_arr_uncertain) == 0:
            break
        # x1 max and x2 min
        x1_max = np.maximum(x1[idx_to_keep], x1[idx_arr_uncertain])
        x2_min = np.minimum(x2[idx_to_keep], x2[idx_arr_uncertain])
        # y1 max and y2 min
        y1_max = np.maximum(y1[idx_to_keep], y1[idx_arr_uncertain])
        y2_min = np.minimum(y2[idx_to_keep], y2[idx_arr_uncertain])
        # intersected x
        xx = x2_min - x1_max + 1
        xx = np.maximum(xx, 0)
        # intersected y
        yy = y2_min - y1_max + 1
        yy = np.maximum(yy, 0)
        # intersected area
        I = xx * yy
        # union area
        U = areas[idx_to_keep] + areas[idx_arr_uncertain] - I
        # IOU: intersection over union
        IOU = I / (U + EPS)

        #########################################################
        # decide which to remove and remove it
        idx_arr_to_remove = idx_arr_uncertain[IOU >= thresh]
        idx_arr_uncertain = idx_arr_uncertain[IOU < thresh]
        print('after')
        print('idx_to_keep', idx_arr_to_keep)
        print('idx_arr_to_remove', idx_arr_to_remove)
        print('idx_arr_uncertain', idx_arr_uncertain)

        #########################################################
        # plot condition of after
        if is_show_demo:
            draw_boxes_groups(boxes, idx_arr_to_keep, idx_arr_to_remove, idx_arr_uncertain, f'#{no_of_action} after')
            plt.show()

    # final
    if is_show_demo:
        plt.figure(figsize=[14, 6])
        spr = 1
        spc = 2
        spn = 0
        idx_arr_removed = np.delete(np.arange(len(boxes)), idx_arr_to_keep)
        draw_boxes_groups(boxes, idx_arr_to_keep, idx_arr_removed, [], f'#{no_of_action} final')
        draw_boxes_groups(boxes, idx_arr_to_keep, [], [], f'result of NMS')
        plt.show()

    # return kept boxes
    return boxes[idx_arr_to_keep]


if '__main__' == __name__:
    ###################################################################################
    # 数据
    print('提示：这个数据第一步没有IOU大于thresh 0.7的，所以第一步没有效果。')
    boxes = np.array([
        [100, 100, 210, 210, 0.72],  # 0
        [280, 290, 420, 420, 0.8],  # 1
        [220, 220, 320, 330, 0.92],  # 2
        [105, 90, 220, 210, 0.71],  # 3
        [230, 240, 325, 330, 0.81],  # 4
        [305, 300, 420, 420, 0.9],  # 5
        [215, 225, 305, 328, 0.6],  # 6
        [150, 260, 290, 400, 0.99],  # 7
        [102, 108, 208, 208, 0.72],  # 8
    ])  # 9个框

    ###################################################################################
    # 纯粹为了结果的调用
    print('抑制前')
    print(boxes)
    nms_result = nms(boxes, is_show_demo=False)  # 不画图的调用方式
    print('抑制后')
    print(nms_result)

    ###################################################################################
    # 为了看演示的调用
    print('开始演示，陆续观察和关闭绘图窗口以继续……')
    nms_result = nms(boxes, is_show_demo=True)  # 画图的调用方式
    print('抑制后')
    print(nms_result)
